from playdoh03 import *
from base import *

import numpy
from numpy import inf,  zeros, tile, nonzero,amin, maximum, minimum,ceil,floor,argsort,array,argmax,sort,mod,cumsum,sum,arange\
,linspace,ones,matrix,log,sqrt,eye,diag,dot,less,exp,triu
from numpy.linalg import eigh,norm
from numpy.random import rand,randn

class CmaEs(Optimization):
    @staticmethod
    def default_values(**opt_info):
            
        if not 'Xmin' in opt_info.keys():
            opt_info['Xmin'] = -inf
        if not 'Xmax' in opt_info.keys():
            opt_info['Xmax'] = inf
        
        
        return opt_info
    
    def initialize(self):
        """
        Initializes the optimization algorithm
        Can use this variable :
            * self.opt_info = dict(algorithm-specific optimization parameters) # static
            self.this_machine_index
        """

        self.sigma = 0.5; 
        self.dim = len(self.params)
        self.N = self.particles
        # Strategy parameter setting: Selection  
        #self.N = 4+floor(3*log(self.dim))                # population size, offspring number
        self.xmean = rand(self.dim,1);    #objective variables initial point
        self.mu = self.N/2                             # number of parents/points for recombination
        self.weights = log(self.mu+1/2)-log(arange(1,self.mu+1))          # muXone array for weighted recombination
        self.mu = floor(self.mu)        
        self.weights = self.weights/sum(self.weights)            # normalize recombination self.weights array
        self.mueff=sum(self.weights)**2/sum(self.weights**2)      # variance-effectiveness of sum w_i x_i
        
        # Strategy parameter setting: Adaptation
        self.cc = (4+self.mueff/self.dim) / (self.dim+4 + 2*self.mueff/self.dim)     # time constant for cumulation for C
        self.cs = (self.mueff+2) / (self.dim+self.mueff+5)             # t-const for cumulation for self.sigma control
        self.c1 = 2 / ((self.dim+1.3)**2+self.mueff)               # learning rate for rank-one update of C
        self.cmu = 2 * (self.mueff-2+1/self.mueff) / ((self.dim+2)**2+self.mueff)  # and for rank-mu update
        self.damps = 1 + 2*max(0, sqrt((self.mueff-1)/(self.dim+1))-1) + self.cs # damping for self.sigma 
                                                        # usually close to 1
        # Initialize dynamic (internal) strategy parameters and constants
        self.pc = zeros((self.dim,1)) 
        self.ps =zeros((self.dim,1))       # evolution paths for C and self.sigma
        self.B = eye(self.dim,self.dim)                            # B defines the coordinate system
        self.D = ones((self.dim))                          # diagonal self.D defines the scaling
        self.C = dot(dot(self.B ,diag(self.D**2)) , self.B.T)                # covariance matrix self.C
        self.invsqrtC = matrix(self.B * diag(self.D**-1) * self.B.T )         # self.C^-1/2 
        self.eigeneval = 0                            # track update of B and self.D
        self.chiN=self.dim**0.5*(1-1/(4*self.dim)+1/(21*self.dim**2))        # expectation of 
        self.counteval=self.N                                
        
        
        Xmin = self.opt_info['Xmin']
        Xmax = self.opt_info['Xmax']
        
        if Xmin is None:
            Xmin = -inf*ones(self.dim)
        if Xmax is None:
            Xmax = inf*ones(self.dim)
        self.Xmin = tile(Xmin.reshape((-1,1)), (1, self.N))
        self.Xmax = tile(Xmax.reshape((-1,1)), (1, self.N))

        
        self.results['best_fits']=list()





    def iterate(self):
        """
        Performs the iteration step of the algorithm.
        The following variables are available :
            * self.param_values = initial parameters # current param values
            * self.X = corresponding initial matrix # current matrix
            * self.fitness = fitness values of initial parameters # current fitness
            * self.iteration = 0 # current iteration
        The variable self.results must be filled by this method.
        The following methods can be used :
            * self.contact : contacts other machines 
            * self.get_fitness(matrix) computes the fitness of the given matrix
            * self.this_machine, other_machines, all_machines in self.contact()
        """
               
        # resort the population from best to worst
         
      
        # Sort by fitness and compute weighted mean into xmean
        self.indices_Population_sorted=argsort(self.fitness)[::-1]

        self.fitness=self.fitness[self.indices_Population_sorted]
        # compute new fitness
                # Generate and evaluate self.N offspring

        self.xold=self.xmean
        self.xmean=dot(self.X[:,self.indices_Population_sorted[0:self.mu]],self.weights) 
        self.xmean=self.xmean.reshape(len(self.xmean),1)
        ind_best=argmax(self.fitness)
        
        
        
        #      [arfitness, arindex] = sort(arfitness); % minimization
        #      xold = xmean;
        #      xmean = arx(:,arindex(1:self.mu))*self.weights;   % recombination, new mean value
        
        # self.Cumulation: Update evolution paths
        
        self.ps = (1-self.cs)*self.ps+ sqrt(self.cs*(2-self.cs)*self.mueff) * dot(self.invsqrtC ,(self.xmean-self.xold)) / self.sigma
        if less(norm(self.ps)/sqrt(1-(1-self.cs)**(2*self.counteval/self.N))/self.chiN , 1.4 + 2/(self.dim+1)):
            hsig=1
        else:
            hsig=0
        
       
        self.pc = (1-self.cc)*self.pc + hsig * sqrt(self.cc*(2-self.cc)*self.mueff) * (self.xmean-self.xold) / self.sigma

        # Adapt covariance matrix self.C
        artmp = (1./self.sigma) * (self.X[:,self.indices_Population_sorted[0:self.mu]]-tile(self.xold,self.mu))
                                   
        self.C = (1-self.c1-self.cmu) * self.C + self.c1 * (dot(self.pc,self.pc.T)+ (1-hsig) * self.cc*(2-self.cc) * self.C)+ self.cmu * dot(dot(artmp , diag(self.weights)) ,artmp.T)
        
        
        # Adapt step size self.sigma
        self.sigma = self.sigma * exp((self.cs/self.damps)*(norm(self.ps)/self.chiN - 1)); 
    
        # self.Decomposition of self.C into B*diag(self.D.^2)*B' (diagonalization)
        #if self.counteval - self.eigeneval > self.N/(self.c1+self.cmu)/self.dim/10:  #to achieve O(self.dim^2)
        self.eigeneval = self.counteval;
        self.C = triu(self.C) + triu(self.C,1).T; # enforce symmetry
        [self.D,self.B] = eigh(self.C);           # eigen decomposition, B==normalized eigenvectors

        self.D = sqrt(self.D);        # self.D is a vector of standard deviations now
        self.invsqrtC = self.B * diag(self.D**-1) * self.B.T

        # Generate and evaluate self.N offspring
        for k in range(self.N):
            self.X[:,k] = self.xmean.T + self.sigma * dot(self.B , (self.D * randn(self.dim))) # m + sig * Normal(0,self.C) 

#        # compute new fitness
        self.get_fitness()
#        self.xmean=dot(self.X[:,self.indices_Population_sorted[0:self.mu]],self.weights) 
#        self.xmean=self.xmean.reshape(len(self.xmean),1)
#        ind_best=argmax(self.fitness)
        
        self.results['best_fit_local'] = self.fitness[ind_best]
        self.results['best_pos_local'] = self.X[:,ind_best]
        self.results['best_fits'].append(self.fitness[ind_best])
        
        self.counteval = self.counteval+self.N;
    
    @staticmethod
    def finalize(results_machines):
        results_final=dict()
        if len(results_machines)==1:
            results_final['best_fit'] = results_machines[0]['best_fit_local']
            results_final['best_pos'] = results_machines[0]['best_pos_local']
            results_final['best_fits']= array(results_machines[0]['best_fits'])
        else:
            best_fits_matrix=zeros((len(results_machines),len(results_machines[0]['best_fits'])))
            best_pos_matrix=zeros((results_machines[0]['best_pos_local'].shape[0],len(results_machines)))
            best_fit_vector=zeros((len(results_machines)))
            for machine_idx in range(len(results_machines)):
                best_fits_matrix[machine_idx,:]=array(results_machines[machine_idx]['best_fits'])
                best_pos_matrix[:,machine_idx]=results_machines[machine_idx]['best_pos_local']
                best_fit_vector[machine_idx]= results_machines[machine_idx]['best_fit_local']
                
            indice_best=argmax(best_fit_vector)
            results_final['best_fit'] = best_fit_vector[indice_best]
            results_final['best_pos'] = best_pos_matrix[:,indice_best]
            results_final['best_fits']= amin(best_fits_matrix,axis=0)
            
        return results_final
